using LinearAlgebra
using KrylovKit

println()
println("Starting echo...")
println()

# matched from experiment
tau_arr = [9.999999999999999e-6, 4.0625e-5, 7.125e-5, 0.000101875, 0.0001325, 0.000163125, 0.00019375, 0.000224375, 0.00025499999999999996, 0.00028562499999999996, 0.00031624999999999996, 0.00034687499999999996, 0.00037749999999999996, 0.00040812499999999996, 0.00043874999999999996, 0.00046937499999999996, 0.0005]

function ggms(d::Int64)::Vector{}

    ggms = Vector{}();
    # first (d^2-d)/2 elements are the "R" operators
    for aa = 1:d
        bsa = spzeros(d,1); bsa[aa,1] = 1;
        for bb = (aa+1):d
            bsb = spzeros(d,1); bsb[bb,1] = 1;
            push!(ggms, (bsb*bsa'+bsa*bsb')/sqrt(2));
        end
    end

    # second (d^2-d)/2 elements are the "L" operators
    for aa = 1:d
        bsa = spzeros(d,1); bsa[aa,1] = 1;
        for bb = 1:(aa-1)
            bsb = spzeros(d,1); bsb[bb,1] = 1;
            push!(ggms, (bsb*bsa'-bsa*bsb')*(-1im)/sqrt(2) );
        end
    end

    # third, the d-1 cartan generators  (Cartan-Weyl basis)
    for mm = 1:(d-1)
          delem = sqrt(1/(mm*(mm+1))).*[ones(1,mm) -mm zeros(1,d-mm-1)];
          push!(ggms, sparse(d:-1:1, d:-1:1, vec(delem)););
      end

      push!(ggms, sparse(I,d,d)/sqrt(d));

    return ggms

end

p0_arr = zeros(length(tau_arr))
p1_arr = zeros(length(tau_arr))

# Making things more efficient using matrix exponential outside
dtau = tau_arr[2] - tau_arr[1]
@assert isapprox(tau_arr, tau_arr[1]:dtau:(tau_arr[end]+dtau/2))
@time evol_tau1_echo = exp(liouville_free * tau_arr[1]/2)
@time evol_dtau_echo = exp(liouville_free * dtau/2)
println("Computing evolution...")
flush(stdout)
@time for (idx_tau, tau) in enumerate(tau_arr)
    p0_tmp = zeros(1)
    p1_tmp = zeros(1)
    evol_tau = evol_dtau_echo^(idx_tau-1)* evol_tau1_echo
    rho0 = zeros(ComplexF64, size(H_x)...)
    for ho_level in 1:num_lvls
        psi0 = zeros(ComplexF64, size(H_x, 1));
        psi0[ho_level] = 1;
        state_probability = 1 / partition_sum * exp(-beta_w_units * energies_N0[ho_level])

        rho0 .+= state_probability .* psi0 .* psi0'
    end
    @assert tr(rho0) ≈ 1
    psit = tr.(GG .* [rho0])

    psit = evol_x_2 * psit
    psit = evol_tau * psit
    psit = evol_y * psit
    psit = evol_tau * psit
    psit = evol_x_2 * psit
    p0 = p0_op_vec'psit
    p1 = p1_op_vec'psit

    p0_arr[idx_tau] = p0
    p1_arr[idx_tau] = p1
end

flush(stdout)

output_df = DataFrame(
    time=tau_arr,
    p0=p0_arr,
    p1=p1_arr,
)

mkpath(save_dir)
ofn = "$save_dir/echo-$pstring.csv"
CSV.write(ofn, output_df)
